// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// A command line executable that generates a bunch of valid IPC files
// containing example record batches.  Those are used as fuzzing seeds
// to make fuzzing more efficient.

#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

#include <nebula/io/file.h>
#include <nebula/io/memory.h>
#include <nebula/ipc/json_simple.h>
#include <tests/ipc/test_common.h>
#include <nebula/ipc/writer.h>
#include <nebula/core/record_batch.h>
#include <turbo/files/file_util.h>
#include <nebula/testing/extension_type.h>
#include <nebula/util/key_value_metadata.h>

namespace nebula {
namespace ipc {
using json::array_from_json;

turbo::Result<std::shared_ptr<RecordBatch>> MakeExtensionBatch() {
  auto array = ExampleUuid();
  auto md = key_value_metadata({"key1", "key2"}, {"value1", ""});
  auto schema = ::nebula::schema({field("f0", array->type())}, md);
  return RecordBatch::create(schema, array->length(), {array});
}

turbo::Result<std::shared_ptr<RecordBatch>> MakeMapBatch() {
  std::shared_ptr<Array> array;
  const char* json_input = R"(
[
    [[0, 1], [1, 1], [2, 2], [3, 3], [4, 5], [5, 8]],
    null,
    [[0, null], [1, null], [2, 0], [3, 1], [4, null], [5, 2]],
    []
  ]
)";
  TURBO_MOVE_OR_RAISE(array, array_from_json(map(int16(), int32()), json_input));
  auto schema = ::nebula::schema({field("f0", array->type())});
  return RecordBatch::create(schema, array->length(), {array});
}

turbo::Result<std::vector<std::shared_ptr<RecordBatch>>> Batches() {
  std::vector<std::shared_ptr<RecordBatch>> batches;
  std::shared_ptr<RecordBatch> batch;
  std::shared_ptr<Array> array;

  TURBO_RETURN_NOT_OK(test::MakeNullRecordBatch(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeListRecordBatch(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeListViewRecordBatch(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeDictionary(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeTimestamps(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeFWBinary(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeStruct(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeUnion(&batch));
  batches.push_back(batch);
  TURBO_RETURN_NOT_OK(test::MakeFixedSizeListRecordBatch(&batch));
  batches.push_back(batch);
  TURBO_MOVE_OR_RAISE(batch, MakeExtensionBatch());
  batches.push_back(batch);
  TURBO_MOVE_OR_RAISE(batch, MakeMapBatch());
  batches.push_back(batch);

  return batches;
}

turbo::Result<std::shared_ptr<Buffer>> serialize_record_batch(
    const std::shared_ptr<RecordBatch>& batch, bool is_stream_format) {
  TURBO_MOVE_OR_RAISE(auto sink, io::BufferOutputStream::create(1024));
  std::shared_ptr<RecordBatchWriter> writer;
  if (is_stream_format) {
    TURBO_MOVE_OR_RAISE(writer, make_stream_writer(sink, batch->schema()));
  } else {
    TURBO_MOVE_OR_RAISE(writer, make_file_writer(sink, batch->schema()));
  }
  TURBO_RETURN_NOT_OK(writer->write_record_batch(*batch));
  TURBO_RETURN_NOT_OK(writer->close());
  return sink->finish();
}

turbo::Status DoMain(bool is_stream_format, const std::string& out_dir) {
  TURBO_MOVE_OR_RAISE(auto dir_fn, turbo::FilePath::from_string(out_dir));
  TURBO_RETURN_NOT_OK(turbo::create_directory(dir_fn));

  int sample_num = 1;
  auto sample_name = [&]() -> std::string {
    return "batch-" + std::to_string(sample_num++);
  };

  TURBO_MOVE_OR_RAISE(auto batches, Batches());

  for (const auto& batch : batches) {
    TURBO_RETURN_NOT_OK(batch->validate_full());
    TURBO_MOVE_OR_RAISE(auto buf, serialize_record_batch(batch, is_stream_format));
    auto sample_fn =  dir_fn/sample_name();
    std::cerr << sample_fn.string() << std::endl;
    TURBO_MOVE_OR_RAISE(auto file, io::FileOutputStream::open(sample_fn.string()));
    TURBO_RETURN_NOT_OK(file->write(buf));
    TURBO_RETURN_NOT_OK(file->close());
  }
  return turbo::OkStatus();
}

TURBO_NORETURN void Usage() {
  std::cerr << "Usage: nebula-ipc-generate-fuzz-corpus "
            << "[-stream|-file] <output directory>" << std::endl;
  std::exit(2);
}

int Main(int argc, char** argv) {
  if (argc != 3) {
    Usage();
  }
  auto opt = std::string(argv[1]);
  if (opt != "-stream" && opt != "-file") {
    Usage();
  }
  auto out_dir = std::string(argv[2]);

  turbo::Status st = DoMain(opt == "-stream", out_dir);
  if (!st.ok()) {
    std::cerr << st.to_string() << std::endl;
    return 1;
  }
  return 0;
}

}  // namespace ipc
}  // namespace nebula

int main(int argc, char** argv) { return nebula::ipc::Main(argc, argv); }
